package drds.data_propagate.driver;

import drds.data_propagate.driver.packets.HeaderPacket;
import drds.data_propagate.driver.packets.client.ClientAuthenticationPacket;
import drds.data_propagate.driver.packets.client.command_packet.AuthSwitchResponsePacket;
import drds.data_propagate.driver.packets.client.command_packet.QuitPacket;
import drds.data_propagate.driver.packets.server.ErrorPacket;
import drds.data_propagate.driver.packets.server.HandshakeInitializationPacket;
import drds.data_propagate.driver.packets.server.Reply323Packet;
import drds.data_propagate.driver.packets.server.command_packet.AuthSwitchRequestMoreData;
import drds.data_propagate.driver.packets.server.command_packet.AuthSwitchRequestPacket;
import drds.data_propagate.driver.socket.SocketChannel;
import drds.data_propagate.driver.socket.SocketChannelPool;
import drds.data_propagate.driver.utils.MSC;
import drds.data_propagate.driver.utils.MySQLPasswordEncrypter;
import drds.data_propagate.driver.utils.PacketManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.security.DigestException;
import java.security.NoSuchAlgorithmException;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 基于mysql socket协议的链接实现
 */
public class Connector {

    public static final int timeout = 5 * 1000;                                     // 5s
    private static final Logger logger = LoggerFactory.getLogger(Connector.class);
    private InetSocketAddress address;
    private String username;
    private String password;
    private byte charsetNumber = 33;
    private String defaultSchema;
    private int soTimeout = 30 * 1000;
    private int connTimeout = 5 * 1000;
    private int receiveBufferSize = 16 * 1024;
    private int sendBufferSize = 16 * 1024;
    private SocketChannel socketChannel;
    private volatile boolean dumping = false;
    // mysql connectinnId
    private long connectionId = -1;
    private AtomicBoolean connected = new AtomicBoolean(false);

    public Connector() {
    }

    public Connector(InetSocketAddress address, String username, String password) {
        String addr = address.getHostString();
        int port = address.getPort();
        this.address = new InetSocketAddress(addr, port);

        this.username = username;
        this.password = password;
    }

    public Connector(InetSocketAddress address, String username, String password, byte charsetNumber,
                     String defaultSchema) {
        this(address, username, password);

        this.charsetNumber = charsetNumber;
        this.defaultSchema = defaultSchema;
    }

    public void connect() throws IOException {
        if (connected.compareAndSet(false, true)) {
            try {
                socketChannel = SocketChannelPool.open(address);
                logger.info("connect MysqlConnection to {}...", address);
                negotiate(socketChannel);
            } catch (Exception e) {
                disconnect();
                throw new IOException("connect " + this.address + " failure", e);
            }
        } else {
            logger.error("the socketChannel can't be connected twice.");
        }
    }

    public void reconnect() throws IOException {
        disconnect();
        connect();
    }

    public void disconnect() throws IOException {
        if (connected.compareAndSet(true, false)) {
            try {
                if (socketChannel != null) {
                    socketChannel.close();
                }
                logger.info("disConnect MysqlConnection to {}...", address);
            } catch (Exception e) {
                throw new IOException("disconnect " + this.address + " failure", e);
            }

            // 执行一次quit
            if (dumping && connectionId >= 0) {
                Connector connector = null;
                try {
                    connector = this.fork();
                    connector.connect();
                    UpdateExecutor updateExecutor = new UpdateExecutor(connector);
                    updateExecutor.update("KILL CONNECTION " + connectionId);
                } catch (Exception e) {
                    // 忽略具体异常
                    logger.info("KILL DUMP " + connectionId + " failure", e);
                } finally {
                    if (connector != null) {
                        connector.disconnect();
                    }
                }

                dumping = false;
            }
        } else {
            logger.info("the socketChannel {} is not connected", this.address);
        }
    }

    public boolean isConnected() {
        return this.socketChannel != null && this.socketChannel.isConnected();
    }

    public Connector fork() {
        Connector connector = new Connector();
        connector.setCharsetNumber(getCharsetNumber());
        connector.setDefaultSchema(getDefaultSchema());
        connector.setAddress(getAddress());
        connector.setPassword(password);
        connector.setUsername(getUsername());
        connector.setReceiveBufferSize(getReceiveBufferSize());
        connector.setSendBufferSize(getSendBufferSize());
        connector.setSoTimeout(getSoTimeout());
        connector.setConnTimeout(connTimeout);
        return connector;
    }

    public void quit() throws IOException {
        QuitPacket quit = new QuitPacket();
        byte[] cmdBody = quit.toBytes();

        HeaderPacket headerPacket = new HeaderPacket();
        headerPacket.setPacketBodyLength(cmdBody.length);
        headerPacket.setPacketSequenceNumber((byte) 0x00);
        PacketManager.writePackets(socketChannel, headerPacket.toBytes(), cmdBody);
    }

    private void negotiate(SocketChannel socketChannel) throws IOException {
        // https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol
        HeaderPacket headerPacket = PacketManager.readHeader(socketChannel, 4, timeout);
        byte[] body = PacketManager.readBytes(socketChannel, headerPacket.getPacketBodyLength(), timeout);
        if (body[0] < 0) {// check field_count
            if (body[0] == -1) {
                ErrorPacket error = new ErrorPacket();
                error.fromBytes(body);
                throw new IOException("handshake exception:\n" + error.toString());
            } else if (body[0] == -2) {
                throw new IOException("Unexpected EOF packet at handshake phase.");
            } else {
                throw new IOException("unpexpected packet with field_count=" + body[0]);
            }
        }
        HandshakeInitializationPacket handshakeInitializationPacket = new HandshakeInitializationPacket();
        handshakeInitializationPacket.fromBytes(body);
        if (handshakeInitializationPacket.protocolVersion != MSC.DEFAULT_PROTOCOL_VERSION) {
            // HandshakeV9
            auth323(socketChannel, (byte) (headerPacket.getPacketSequenceNumber() + 1), handshakeInitializationPacket.seed);
            return;
        }

        connectionId = handshakeInitializationPacket.threadId; // 记录一下connection
        logger.info("handshake initialization packet received, prepare the client authentication packet to send");
        ClientAuthenticationPacket clientAuthenticationPacket = new ClientAuthenticationPacket();
        clientAuthenticationPacket.setCharsetNumber(charsetNumber);

        clientAuthenticationPacket.setUsername(username);
        clientAuthenticationPacket.setPassword(password);
        clientAuthenticationPacket.setServerCapabilities(handshakeInitializationPacket.serverCapabilities);
        clientAuthenticationPacket.setDatabaseName(defaultSchema);
        clientAuthenticationPacket.setScrumbleBuff(joinAndCreateScrumbleBuff(handshakeInitializationPacket));
        clientAuthenticationPacket.setAuthPluginName("mysql_native_password".getBytes());

        byte[] clientAuthPkgBody = clientAuthenticationPacket.toBytes();
        HeaderPacket headerPacket2 = new HeaderPacket();
        headerPacket2.setPacketBodyLength(clientAuthPkgBody.length);
        headerPacket2.setPacketSequenceNumber((byte) (headerPacket.getPacketSequenceNumber() + 1));

        PacketManager.writePackets(socketChannel, headerPacket2.toBytes(), clientAuthPkgBody);
        logger.info("client authentication packet is sent out.");

        // check auth result
        headerPacket = null;
        headerPacket = PacketManager.readHeader(socketChannel, 4);
        body = null;
        body = PacketManager.readBytes(socketChannel, headerPacket.getPacketBodyLength(), timeout);
        assert body != null;
        byte marker = body[0];
        if (marker == -2 || marker == 1) {
            byte[] authData = null;
            String pluginName = null;
            if (marker == 1) {
                AuthSwitchRequestMoreData packet = new AuthSwitchRequestMoreData();
                packet.fromBytes(body);
                authData = packet.authData;
            } else {
                AuthSwitchRequestPacket packet = new AuthSwitchRequestPacket();
                packet.fromBytes(body);
                authData = packet.authData;
                pluginName = packet.authName;
            }

            boolean isSha2Password = false;
            byte[] encryptedPassword = null;
            if (pluginName != null && "mysql_native_password".equals(pluginName)) {
                try {
                    encryptedPassword = MySQLPasswordEncrypter.scramble411(getPassword().getBytes(), authData);
                } catch (NoSuchAlgorithmException e) {
                    throw new RuntimeException("can't encrypt password that will be sent to MySQL server.", e);
                }
            } else if (pluginName != null && "caching_sha2_password".equals(pluginName)) {
                isSha2Password = true;
                try {
                    encryptedPassword = MySQLPasswordEncrypter.scrambleCachingSha2(getPassword().getBytes(), authData);
                } catch (DigestException e) {
                    throw new RuntimeException("can't encrypt password that will be sent to MySQL server.", e);
                }
            }
            assert encryptedPassword != null;
            AuthSwitchResponsePacket responsePacket = new AuthSwitchResponsePacket();
            responsePacket.authData = encryptedPassword;
            byte[] auth = responsePacket.toBytes();

            headerPacket2 = new HeaderPacket();
            headerPacket2.setPacketBodyLength(auth.length);
            headerPacket2.setPacketSequenceNumber((byte) (headerPacket.getPacketSequenceNumber() + 1));
            PacketManager.writePackets(socketChannel, headerPacket2.toBytes(), auth);
            logger.info("auth switch response packet is sent out.");

            headerPacket = null;
            headerPacket = PacketManager.readHeader(socketChannel, 4);
            body = null;
            body = PacketManager.readBytes(socketChannel, headerPacket.getPacketBodyLength(), timeout);
            assert body != null;
            if (isSha2Password) {
                if (body[0] == 0x01 && body[1] == 0x04) {
                    // password auth failed
                    throw new IOException("caching_sha2_password Auth failed");
                }

                headerPacket = null;
                headerPacket = PacketManager.readHeader(socketChannel, 4);
                body = null;
                body = PacketManager.readBytes(socketChannel, headerPacket.getPacketBodyLength(), timeout);
            }
        }

        if (body[0] < 0) {
            if (body[0] == -1) {
                ErrorPacket err = new ErrorPacket();
                err.fromBytes(body);
                throw new IOException("Error When doing Client Authentication:" + err.toString());
            } else {
                throw new IOException("unpexpected packet with field_count=" + body[0]);
            }
        }
    }

    private void auth323(SocketChannel channel, byte packetSequenceNumber, byte[] seed) throws IOException {
        // auth 323
        Reply323Packet r323 = new Reply323Packet();
        if (password != null && password.length() > 0) {
            r323.seed = MySQLPasswordEncrypter.scramble323(password, new String(seed)).getBytes();
        }
        byte[] b323Body = r323.toBytes();

        HeaderPacket h323 = new HeaderPacket();
        h323.setPacketBodyLength(b323Body.length);
        h323.setPacketSequenceNumber((byte) (packetSequenceNumber + 1));

        PacketManager.writePackets(channel, h323.toBytes(), b323Body);
        logger.info("client 323 authentication packet is sent out.");
        // check auth result
        HeaderPacket header = PacketManager.readHeader(channel, 4);
        byte[] body = PacketManager.readBytes(channel, header.getPacketBodyLength());
        assert body != null;
        switch (body[0]) {
            case 0:
                break;
            case -1:
                ErrorPacket err = new ErrorPacket();
                err.fromBytes(body);
                throw new IOException("Error When doing Client Authentication:" + err.toString());
            default:
                throw new IOException("unpexpected packet with field_count=" + body[0]);
        }
    }

    private byte[] joinAndCreateScrumbleBuff(HandshakeInitializationPacket handshakePacket) throws IOException {
        byte[] dest = new byte[handshakePacket.seed.length + handshakePacket.restOfScrambleBuff.length];
        System.arraycopy(handshakePacket.seed, 0, dest, 0, handshakePacket.seed.length);
        System.arraycopy(handshakePacket.restOfScrambleBuff,
                0,
                dest,
                handshakePacket.seed.length,
                handshakePacket.restOfScrambleBuff.length);
        return dest;
    }

    public InetSocketAddress getAddress() {
        return address;
    }

    public void setAddress(InetSocketAddress address) {
        this.address = address;
    }

    public String getUsername() {
        return username;
    }

    public void setUsername(String username) {
        this.username = username;
    }

    public byte getCharsetNumber() {
        return charsetNumber;
    }

    public void setCharsetNumber(byte charsetNumber) {
        this.charsetNumber = charsetNumber;
    }

    public String getDefaultSchema() {
        return defaultSchema;
    }

    public void setDefaultSchema(String defaultSchema) {
        this.defaultSchema = defaultSchema;
    }

    public int getSoTimeout() {
        return soTimeout;
    }

    public void setSoTimeout(int soTimeout) {
        this.soTimeout = soTimeout;
    }

    public int getReceiveBufferSize() {
        return receiveBufferSize;
    }

    public void setReceiveBufferSize(int receiveBufferSize) {
        this.receiveBufferSize = receiveBufferSize;
    }

    public int getSendBufferSize() {
        return sendBufferSize;
    }

    public void setSendBufferSize(int sendBufferSize) {
        this.sendBufferSize = sendBufferSize;
    }

    public SocketChannel getSocketChannel() {
        return socketChannel;
    }

    public void setSocketChannel(SocketChannel socketChannel) {
        this.socketChannel = socketChannel;
    }

    public long getConnectionId() {
        return connectionId;
    }

    public void setConnectionId(long connectionId) {
        this.connectionId = connectionId;
    }

    public boolean isDumping() {
        return dumping;
    }

    public void setDumping(boolean dumping) {
        this.dumping = dumping;
    }

    public int getConnTimeout() {
        return connTimeout;
    }

    public void setConnTimeout(int connTimeout) {
        this.connTimeout = connTimeout;
    }

    public String getPassword() {
        return password;
    }

    public void setPassword(String password) {
        this.password = password;
    }

}
